# **************************** LICENSE START ***********************************
#
# Copyright 2012 ECMWF. This software is distributed under the terms
# of the Apache License version 2.0. In applying this license, ECMWF does not
# waive the privileges and immunities granted to it by virtue of its status as
# an Intergovernmental Organization or submit itself to any jurisdiction.
#
# ***************************** LICENSE END ************************************

extern gradientb(f:fieldset) "C" inline

/*
   "gradientb" - computes the gradient of a field.

*/




#include <stdio.h>
#include <string.h>
#include <math.h>
#include "macro_api.h"


/*int idx (int X, int Y)
{
    printf ("%d, %d\n", X, Y);
    return ((Y)*x_num + (X));
}*/

/*
    check_data_ok - checks whether the data is in a format we can work with
                  - returns 1 if it is, 0 otherwise
*/

int check_data_ok (grib_handle *gh)
{
    char grid_type [32];
    size_t len = sizeof(grid_type);
    
    grib_get_string (gh, "typeOfGrid", grid_type, &len);
    
    if (strcmp (grid_type, "regular_ll"))
    {
        printf ("Data is in wrong grid type (%s) - it should be 'regular_ll'\n", grid_type);
        return 0;
    }
    

    /* if we got to here, then it was all good */
    
    return 1;
}


void compute_gradient (grib_handle *gh, double *vals_out_u, double *vals_out_v)
{
    #define idx(X,Y) ((Y)*x_num + (X))

    double pi = acos(-1.0);
    double rt = 6371000.0;
    double cb = (rt * 1.5 * pi) / 180.0;
    long x_num, y_num;
    double x_inc, y_inc;
    double *vals;
    int ret;
    size_t len;
    int i, j;

    if (!check_data_ok(gh))
    {
        mci_fail("Data not ok for this function - aborting.");
    }


    GRIB_CHECK(grib_get_long   (gh, "numberOfPointsAlongAParallel", &x_num), 0);
    GRIB_CHECK(grib_get_long   (gh, "numberOfPointsAlongAMeridian", &y_num), 0);
    GRIB_CHECK(grib_get_double (gh, "iDirectionIncrement", &x_inc), 0);
    GRIB_CHECK(grib_get_double (gh, "jDirectionIncrement", &y_inc), 0);
    
    x_inc /= 1000.0; /* increments are stored in millidegrees */
    y_inc /= 1000.0; /* increments are stored in millidegrees */


    /* check that the data is global */

    if ((x_num * x_inc != 360.0) || ((y_num-1) * y_inc != 180.0))
    {
        printf ("Data is not global (%d x %f, %d x %f)\n", x_num, x_inc, y_num, y_inc);
        mci_fail("Data not ok for this function - aborting.");
    }



    len = x_num * y_num;

    vals = (double *) malloc (len * sizeof(double));

    printf ("getting %d elements...\n", len);
    ret = grib_get_double_array(gh,"values",vals, &len);
    printf ("got %d elements...\n", len);

    if (ret == GRIB_SUCCESS)
    {
        /* COMPUTE HORIZONTAL GRADIENT */

        for (i = 0; i < y_num; i++)
        {
            double c = cos((90.0 - (i * y_inc) + x_inc) * pi/180.0);

            if ((fabs(c) < 0.00001))
            {
                c = 0.00001;
            }

            vals_out_u[idx(0,   i)] = (vals[idx(1, i)] - vals[idx(239, i)]) / (2.0 * c * cb);
            vals_out_u[idx(239, i)] = (vals[idx(0, i)] - vals[idx(238, i)]) / (2.0 * c * cb);

            for (j = 1; j < x_num-1; j++)
            {
                vals_out_u[idx(j,i)] = (vals[idx(j+1,i)] - vals[idx(j-1,i)]) / (2.0 * c * cb);
            }
        }


        /* COMPUTE VERTICAL GRADIENT */

        for (i = 0; i < x_num; i++)
        {
            vals_out_v[idx(i,0)]   = 0;
            vals_out_v[idx(i,120)] = 0;
            
            for (j = 1; j < y_num-1; j++)
            {
                vals_out_v[idx(i,j)] = (vals[idx(i,j+1)] - vals[idx(i,j-1)]) / (-2.0 * cb);
            }
        }
    }

    else
    {
      printf(">>> ERROR: grib_get_double_array returned %d\n", ret);
    }
}



int main()
{
    grib_handle* gh_a = NULL;
    grib_handle* gh_u = NULL;
    grib_handle* gh_v = NULL;
    void*   grib_id = NULL;

    int   num_fields   = 0;    /*-- field count --*/
    int    ret_a   = 0;    /*-- function return value --*/
    int    ret_b   = 0;    /*-- function return value --*/
    int    i;
    size_t len_a   = 0;    /*-- number of grid point values --*/
    size_t len_b   = 0;    /*-- number of grid point values --*/
    void*   grib_out_id = NULL; /*-- return GRIB id ptr, fieldset handle */
    double *outvals_u = NULL;        /*-- array for grid point values */
    double *outvals_v = NULL;        /*-- array for grid point values */


    /* load the input grib fields */

    grib_id = mci_get_grib_id_ptr (&num_fields);


    /* loop on the input fields */
    
    for (i = 0; i < num_fields; i++)
    {
        /* get arrays of values from the fields */

        gh_a = mci_load_one_grib   (grib_id);

        grib_get_size(gh_a,"values",&len_a);
        printf("GRID size: %d\n",len_a);

        grib_out_id = mci_new_grib_id_ptr();  /* create return GRIB id ptr */


        /* allocate memory for the result arrays */

        outvals_u = (double *) malloc (len_a * sizeof (double));
        outvals_v = (double *) malloc (len_a * sizeof (double));


        /* compute the gradient fields */

        compute_gradient (gh_a, outvals_u, outvals_v);


        /* create new fields using these results */

        gh_u = grib_handle_clone (gh_a);
        gh_v = grib_handle_clone (gh_a);


        /* set their values using the arrays we've computed */

        grib_set_double_array(gh_u, "values", outvals_u, len_a);
        grib_set_double_array(gh_v, "values", outvals_v, len_a);


        /* set their parameters to U and V */

        grib_set_long (gh_u, "indicatorOfParameter", 131);
        grib_set_long (gh_v, "indicatorOfParameter", 132);


        /* save them as the result of this inline program */

        mci_save_grib   (grib_out_id, gh_u);
        mci_save_grib   (grib_out_id, gh_v);
        mci_return_grib (grib_out_id);

        if (outvals_u != NULL) free(outvals_u);
        if (outvals_v != NULL) free(outvals_v);

        grib_handle_delete(gh_a);
        grib_handle_delete(gh_u);
        grib_handle_delete(gh_v);
    }

    free(grib_id);
    
    return 0;
}


end inline




# set the area we wish to retrieve data from
#          N,   W,  S,  E

area_xx = [70, -45, 10, 85]


# Retrieve the specific humidity
q = retrieve(
		date	:	-1,
		param	:	"q",
		level	:	700,
		grid	:	[1.5,1.5]
		)

# Get the u and v components of the wind
u = retrieve(
		date	:	-1,
		param	:	"u",
		level	:	700,
		area	:	area_xx,
		grid	:	[1.5,1.5]
		)
v = retrieve(
		date	:	-1,
		param	:	"v",
		level	:	700,
		area	:	area_xx,
		grid	:	[1.5,1.5]
		)


# Compute the gradient of Q
q = gradientb(q)


# Extract the area we are calculating on
q = read ( area : area_xx, data : q)



# Compute the advection of Q
a = q[1]*u + q[2]*v
a = -a * (10 ^ 8) # units will be 10e-8 (kg/kg)/sec


# Plot positive advection in blue, negative in red
contour_common = (
		contour_level_selection_type    :	"interval",
		contour_interval                :	3,
		contour_label                   :	"on",
		contour_label_height            :	0.25,	
		contour_highlight               :	"off",
		contour_hilo                    :	"on",
		contour_hilo_type               :	"number",
		contour_hilo_format             :	"F5.1",
		contour_hilo_height            :	0.3
		)

cont_n = mcont(
		contour_common,
		contour_max_level       :	-0.0001,
		contour_line_colour     :	"red",
		contour_label_colour    :	"red",
		contour_lo_colour       :	"red"
		)

cont_p = mcont(
		contour_common,
		contour_min_level       :	0.0001,
		contour_line_colour     :	"blue",
		contour_label_colour    :	"blue",
		contour_hi_colour       :	"blue"
		)

# A plot window
acoast = mcoast(
		map_coastline_resolution        :	"low",
		map_grid_longitude_increment    :	10,
		map_coastline_land_shade        :	"on",
		map_coastline_land_shade_colour :	"cream"
		)

ps_atlantic = geoview(
		map_projection  :	"polar_stereographic",
		area            :	[30,-25,50,65],
		coastlines      :	acoast
		)

page = plot_page(
		view            :	ps_atlantic
		)

dw = plot_superpage(
		custom_width    :	29.7,
		custom_height   :	21,
		pages           :	page
		)



# Now plot the result

plot(dw,a,cont_p,cont_n)

